//Util.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Runtime.InteropServices;
using size_t = nuint;

unsafe static class Util
{
	public static bool IsWindows => RuntimeInformation.IsOSPlatform(OSPlatform.Windows);
	public static bool IsLinux => RuntimeInformation.IsOSPlatform(OSPlatform.Linux);
	public static bool IsX64 => RuntimeInformation.ProcessArchitecture == Architecture.X64;

	public static size_t LargePageBytes => 2 * 1024 * 1024;
	public static size_t PageBytes => 4 * 1024;

	// if failed, use CacheAlignAlloc instead
	public static double* TryAllocLargePages(size_t doubleNum)
	{
		double* result = null;
		if (IsWindows) {
			result = (double*)WindowsStuff.AllocPages(doubleNum * sizeof(double), false, true);
		}
		else if (IsLinux) {
			result = (double*)LinuxStuff.AllocPages(doubleNum * sizeof(double), false, true);
		}

		if (result != null)
			Console.WriteLine($"tid:{Thread.CurrentThread.ManagedThreadId}\tLarge page allocated {doubleNum * sizeof(double)} bytes");
		return result;
	}

	public static void FreeLargePages(double* addr, size_t doubleNum)
	{
		if (!(IsWindows || IsLinux))
			throw new PlatformNotSupportedException();

		bool succeed = false;
		if (IsWindows) {
			// Windows requires size to be 0 when freeing large pages
			succeed = WindowsStuff.FreePages(addr, 0);
		}
		if (IsLinux) {
			succeed = LinuxStuff.FreePages(addr, doubleNum * sizeof(double));
		}
		if (!succeed)
			throw new Exception($"FreeLargePages failed. err:0x{Marshal.GetLastWin32Error():x}");
	}

	public static void* AllocExecPages(size_t bytes)
	{
		if (IsWindows)
			return WindowsStuff.AllocPages(bytes, true, false);
		if (IsLinux)
			return LinuxStuff.AllocPages(bytes, true, false);
		throw new PlatformNotSupportedException();
	}

	public static size_t CachelineBytes => 64;

	public static double* CacheAlignAlloc(size_t doubleNum)
	{
		size_t bytesRequired = doubleNum * sizeof(double) + 1;

		byte* heapMem = (byte*)Marshal.AllocHGlobal((IntPtr)(ulong)(bytesRequired + CachelineBytes - 1));
		byte* nextCacheLine = (byte*)RoundUpPow2((size_t)heapMem + 1, CachelineBytes);

		size_t padding = (size_t)(nextCacheLine - heapMem); // padding in [1,64]
		*(nextCacheLine - 1) = (byte)padding;
		return (double*)nextCacheLine;
	}

	public static void CacheAlignFree(double* addr)
	{
		if (addr == null)
			return;
		size_t padding = *((byte*)addr - 1);
		byte* heapMem = (byte*)addr - padding;
		Marshal.FreeHGlobal((IntPtr)heapMem);
	}

	public static size_t RoundUpPow2(size_t n, size_t pow2)
	{
		return (n + pow2 - 1) & ~(pow2 - 1);
	}

	public static size_t RoundDownPow2(size_t n, size_t pow2)
	{
		return n & ~(pow2 - 1);
	}
	
	static int HexToInt32(char hex)
	{
		if (hex >= '0' && hex <= '9')
			return hex - '0';
		if (hex >= 'A' && hex <= 'F')
			return hex - 'A' + 10;
		if (hex >= 'a' && hex <= 'f')
			return hex - 'a' + 10;
		return -1;
	}

	public static byte[] ConvertBytes(string[] arr)
	{
		var bytes = new List<byte>();
		int estSize = (arr.Select(line => line.Length).Sum() + 2) / 3;
		bytes.Capacity = estSize;

		bool hexEven = true;
		int byteValue = 0;
		foreach (string line in arr) {
			foreach (char c in line) {
				int value = HexToInt32(c);
				if (value < 0)
					continue;
				if (hexEven) {
					byteValue = value;
				}
				else {
					byteValue = byteValue * 16 + value;
					bytes.Add((byte)byteValue);
				}
				hexEven = !hexEven;
			}
		}
		return bytes.ToArray();
	}
}

unsafe static class WindowsStuff
{
	[DllImport("Kernel32.dll", SetLastError = true)]
	public extern static void* VirtualAlloc(void* addr, UIntPtr size, uint allocType, uint protect);

	[DllImport("Kernel32.dll", SetLastError = true)]
	extern static int VirtualFree(void* addr, UIntPtr size, uint freeType);

	const uint MEM_COMMIT = 0x00001000;
	const uint MEM_RESERVE = 0x00002000;
	const uint MEM_RELEASE = 0x00008000;
	const uint MEM_LARGE_PAGES = 0x20000000;
	const uint PAGE_READWRITE = 0x04;
	const uint PAGE_EXECUTE_READWRITE = 0x40;

	public static void* AllocPages(size_t size, bool exec, bool largePage)
	{
		if (largePage && !lockMemoryEnabled)
			return null;
		return VirtualAlloc(null, size,
			MEM_RESERVE | MEM_COMMIT | (largePage ? MEM_LARGE_PAGES : 0),
			exec ? PAGE_EXECUTE_READWRITE : PAGE_READWRITE);
	}

	public static bool FreePages(void* addr, size_t size)
	{
		return 0 != VirtualFree(addr, size, MEM_RELEASE);
	}

	static bool lockMemoryEnabled;
	
	static WindowsStuff()
	{
		lockMemoryEnabled = EnableLockMemory();
	}

	[DllImport("Advapi32.dll", SetLastError = true)]
	extern static int OpenProcessToken(IntPtr processHandle, uint desiredAcess, out IntPtr tokenHandle);

	[DllImport("Advapi32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
	extern static int LookupPrivilegeValueW(string systemName, string name, out Luid tokenHandle);

	[DllImport("Advapi32.dll", SetLastError = true)]
	extern static int AdjustTokenPrivileges(IntPtr tokenHandle, int disableAllPrivileges, in TokenPrivileges newState,
											uint bufferLength, TokenPrivileges* previousState, uint* returnLength);

	const uint TOKEN_QUERY = 0x0008;
	const uint TOKEN_ADJUST_PRIVILEGES = 0x0020;
	const uint SE_PRIVILEGE_ENABLED = 0x00000002;
	const string SE_LOCK_MEMORY_NAME = "SeLockMemoryPrivilege";
	const int ERROR_SUCCESS = 0;

	[StructLayout(LayoutKind.Sequential)]
	struct Luid
	{
		public uint LowPart;
		public int HightPart;
	}

	[StructLayout(LayoutKind.Sequential)]
	struct TokenPrivileges
	{
		public uint PrivilegeCount;
		public Luid Luid;
		public uint Attributes;
	}

	static bool EnableLockMemory()
	{
		IntPtr hProc = System.Diagnostics.Process.GetCurrentProcess().Handle;
		IntPtr hToken;
		if (0 == OpenProcessToken(hProc, TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, out hToken)) {
			// fail
			return false;
		}
		TokenPrivileges tp;
		tp.PrivilegeCount = 1;
		tp.Attributes = SE_PRIVILEGE_ENABLED;
		LookupPrivilegeValueW(null, SE_LOCK_MEMORY_NAME, out tp.Luid);

		int success = AdjustTokenPrivileges(hToken, 0, tp, 0, null, null);
		if (success != 0 && Marshal.GetLastWin32Error() == ERROR_SUCCESS) {
			// succeed
			return true;
		}
		return false;
	}
}

unsafe static class LinuxStuff
{
	[DllImport("libc", SetLastError = true)]
	extern static void* mmap(void* addr, UIntPtr length, int prot, int flags, int fd, IntPtr offset);

	[DllImport("libc", SetLastError = true)]
	extern static int munmap(void* addr, UIntPtr length);

	const int PROT_READ = 0x1;
	const int PROT_WRITE = 0x2;
	const int PROT_EXEC = 0x4;
	const int MAP_PRIVATE = 0x02;
	const int MAP_ANONYMOUS = 0x20;
	const int MAP_HUGETLB = 0x40000;
	const int MAP_HUGE_SHIFT = 26;
	const int MAP_HUGE_2MB = 21 << MAP_HUGE_SHIFT;

	// no throw
	public static void* AllocPages(size_t size, bool exec, bool largePage)
	{
		void* result = mmap(null, size,
			PROT_READ | PROT_WRITE | (exec ? PROT_EXEC : 0),
			MAP_PRIVATE | MAP_ANONYMOUS | (largePage ? (MAP_HUGETLB | MAP_HUGE_2MB) : 0),
			-1, (IntPtr)0);
		if (result == (void*)(long)-1)
			return null;
		else
			return result;
	}

	// no throw
	public static bool FreePages(void* addr, size_t size)
	{
		return 0 == munmap(addr, size);
	}
}
